/**
 * 
 */
package org.perfect.batch.configuration;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;

import javax.sql.DataSource;

import org.apache.shardingsphere.api.hint.HintManager;
import org.perfect.batch.BatchConstants;
import org.perfect.batch.partition.DataSourcePartitioner;
import org.perfect.batch.step.PartitionMasterStep;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.batch.core.Job;
import org.springframework.batch.core.Step;
import org.springframework.batch.core.configuration.annotation.JobBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepScope;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.batch.item.ItemStreamException;
import org.springframework.batch.item.ItemWriter;
import org.springframework.batch.item.database.JdbcPagingItemReader;
import org.springframework.batch.item.database.support.SqlPagingQueryProviderFactoryBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.RowMapper;

/**
 * @author billmse
 *
 */
@Configuration
public class PartitionJobConfiguration {

	@Autowired
	private JobBuilderFactory jobBuilderFactory;

	@Autowired
	private StepBuilderFactory stepBuilderFactory;
	
	@Autowired
	private PartitionMasterStep partitionMasterStep;
	
	@Autowired
	private DataSourcePartitioner partitioner;
	
	@Autowired
	private DataSource dataSource;
	
	public static final String JOB_NAME = "PartitionJob";
	public static final String MASTER_STEP_NAME = "PartitionMasterStep";
	public static final String SLAVE_STEP_NAME = "PartitionSlaveStep";
	public static final String SLAVE_STEP_READER_NAME = "PartitionSlaveStepReader";
	public static final String SLAVE_STEP_WRITER_NAME = "PartitionSlaveStepWriter";

	private static final int CHUNK_SIZE = 10;
	
	private Logger logger = LoggerFactory.getLogger(getClass());

	@Bean(JOB_NAME)
	public Job partitionJob() {
		return jobBuilderFactory.get(JOB_NAME).start(partitionMasterStep()).build();
	}

	@Bean(MASTER_STEP_NAME)
	public Step partitionMasterStep() {
		return partitionMasterStep.build(MASTER_STEP_NAME, SLAVE_STEP_NAME, partitioner);
	}

	@Bean(SLAVE_STEP_NAME)
	public Step partitionSlaveStep() throws Exception {
		return stepBuilderFactory.get(SLAVE_STEP_NAME).<String, String>chunk(CHUNK_SIZE)
				.reader(partitionSlaveStepReader()).writer(partitionSlaveStepWriter()).build();
	}

	@StepScope
	@Bean(SLAVE_STEP_READER_NAME)
	public PartitionSlaveStepReader partitionSlaveStepReader() throws Exception {
		PartitionSlaveStepReader reader = new PartitionSlaveStepReader();
		reader.setDataSource(dataSource);
		reader.setRowMapper(new RowMapper<String>() {

			@Override
			public String mapRow(ResultSet rs, int rowNum) throws SQLException {
				return rs.getString(1);
			}
			
		});
		SqlPagingQueryProviderFactoryBean sqlFactory = new SqlPagingQueryProviderFactoryBean();
		sqlFactory.setDataSource(dataSource);
		sqlFactory.setSelectClause("*");
		sqlFactory.setFromClause("item");
		sqlFactory.setSortKey("id");
		reader.setQueryProvider(sqlFactory.getObject());
		return reader;
	}
	
	class PartitionSlaveStepReader extends JdbcPagingItemReader<String> {
		
		private static final String LOGIC_TABLE_NAME = "item";
		
		private Logger logger = LoggerFactory.getLogger(getClass());
		
		@Override
		public void open(ExecutionContext executionContext) {
			String logicTableName = executionContext.getString(BatchConstants.BATCH_PARTITION_LOGIC_TABLE_NAME);
			String dataSourceName = executionContext.getString(BatchConstants.BATCH_PARTITION_DATASOURCE_NAME);
			String tableName = executionContext.getString(BatchConstants.BATCH_PARTITION_TABLE_NAME);
			if (LOGIC_TABLE_NAME.equals(logicTableName)) {
				HintManager hintManager = HintManager.getInstance();
				hintManager.addDatabaseShardingValue(logicTableName, dataSourceName.substring(dataSourceName.length() - 1));
				hintManager.addTableShardingValue(logicTableName, tableName.substring(tableName.length() - 1));
				logger.info("{}, {}, {}", logicTableName, dataSourceName, tableName);
			}
			super.open(executionContext);
		}
		
		@Override
		public void close() throws ItemStreamException {
			super.close();
			HintManager.clear();
		}
		
	}

	@StepScope
	@Bean(SLAVE_STEP_WRITER_NAME)
	public PartitionSlaveStepWriter partitionSlaveStepWriter() {
		return new PartitionSlaveStepWriter();
	}
	
	class PartitionSlaveStepWriter implements ItemWriter<String> {

		@Override
		public void write(List<? extends String> items) throws Exception {
			logger.info("write items: {}", items);
		}
		
	}

}
